#include "head.h"
#include "lua_socket.h"

/*
    usage:  connect(ip, port, [,AF_INET]);
    paras:  ip, port, [, AF_INET]
    return: [sockfd]、[nil, errno]
*/
static int l_socket_connect(lua_State *L) {
    const char *ip; 
    int sockfd, port; 
    short sin_family; 
    struct sockaddr_in server;

    ip = luaL_checkstring(L, 1);
    port = luaL_checkinteger(L, 2);
    sin_family = (short)luaL_optinteger(L, 3, AF_INET);

    sockfd = socket(PF_INET, SOCK_STREAM, 0); 
    if (sockfd == -1) {
        lua_pushnil(L);
        lua_pushstring(L, strerror(errno));
        return 2;
    }

    bzero(&server, sizeof(server));
    server.sin_addr.s_addr = inet_addr(ip);
    server.sin_port = htons(port);
    server.sin_family = sin_family; 

    if (connect(sockfd, (struct sockaddr *)&server, sizeof(struct sockaddr)) == -1) {
        lua_pushnil(L);
        lua_pushstring(L, strerror(errno));
        return 2;
    }
    lua_pushinteger(L, sockfd);
    return 1;
}

/*
    usage:  close(sockfd)
    paras:  sockfd
    return: [errno]
*/
static int l_socket_close(lua_State *L) {
    int sockfd = luaL_checkinteger(L, -1);
    if (close(sockfd) == -1) {
        lua_pushstring(L, strerror(errno));
        return 1;
    }
    return 0;
}

/*
    usage:  send(sockfd, "MESSAGE")
    paras:  sockfd, msg
    return: [len]
*/
static int l_socket_send(lua_State *L) {
    size_t sz; 
    int sockfd = luaL_checkinteger(L, 1);
    const char *msg = luaL_checklstring(L, 2, &sz); // 第二个参数大小存给sz

    int len = send(sockfd, msg, (int)sz, 0);
    lua_pushinteger(L, len);
    return 1;
}

/*
    usage:  recv(sockfd, [, size])
    paras:  sockfd, [, size -- default: 1024]
    return: [buffer]、["", "closed"]、[nil, "timeout"]、[nil, errno]
*/
static int l_socket_recv(lua_State *L) {
    int sockfd = luaL_checkinteger(L, 1);
    int sz = luaL_optinteger(L, 2, 1024);
    char buffer[1024];

    int len = recv(sockfd, buffer, sz, 0);
    if (len > 0) {
        lua_pushlstring(L, buffer, len);
        return 1;
    } else if (len == 0) {
        lua_pushstring(L, "");
        lua_pushstring(L, "closed");
        return 2;
    } else {
        if (errno == EAGAIN) {
            lua_pushnil(L);
            lua_pushstring(L, "timeout");
            return 2;
        } else {
            lua_pushnil(L);
            lua_pushstring(L, strerror(errno));
            return 2;
        }
    }
}

/*
    usage:  listen(ip, port, [, AF_INET])
    paras:  ip, port, [, AF_INET]
    return: [listener]
*/
static int l_socket_listen(lua_State *L) {
    static int reuse = -1;
    const char *ip; 
    int port, listener;
    struct sockaddr_in my_addr;
    short sin_family; 

    ip = luaL_checkstring(L, 1);
    port = luaL_checkinteger(L, 2);
    sin_family = (short)luaL_optinteger(L, 3, AF_INET);

    if ((listener = socket(PF_INET, SOCK_STREAM, 0)) == -1) {
        lua_pushnil(L);
        perror("socket");
        return 0;
    }

    bzero(&my_addr, sizeof(my_addr));

    my_addr.sin_addr.s_addr = inet_addr(ip);
    my_addr.sin_port = htons(port);
    my_addr.sin_family = sin_family;

    if (setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, (void *)&reuse, sizeof(int)) == -1) {
        lua_pushnil(L);
        perror("setsockopt");
        return 0;
    }

    if (bind(listener, (struct sockaddr *)&my_addr, sizeof(struct sockaddr)) == -1) {
        lua_pushnil(L);
        perror("bind");
        return 0;
    }

    if (listen(listener, 32) == -1) {
        lua_pushnil(L);
        perror("listen");
        return 0;
    }

    lua_pushinteger(L, listener);
    return 1;
}

/*
    usage:  accept(listener)
    paras:  listener
    return: [sockfd, "ip:port"]、[nil, nil, "timeout"]、[nil, nil, errno]
*/
static int l_socket_accept(lua_State *L) {
    static struct sockaddr_in client; 
    static socklen_t client_len = sizeof(struct sockaddr_in); 

    char *ip;
    int port; 

    int listener = luaL_checkinteger(L, 1);
    int sockfd = accept(listener, (struct sockaddr *)&client, &client_len); 

    if (sockfd < 0) {
        if (errno == EAGAIN) {
            lua_pushnil(L);
            lua_pushnil(L);
            lua_pushstring(L, "timeout");
            return 3;
        } else {
            lua_pushnil(L);
            lua_pushnil(L);
            lua_pushstring(L, strerror(errno));
            return 3;
        }
    }

    ip = inet_ntoa(client.sin_addr);
    port = ntohs(client.sin_port);

    int len = strlen(ip);
    char addr[len + 10];
    memcpy(addr, ip, len);
    addr[len] = ':';
    sprintf(addr + len + 1, "%d", port);

    lua_pushinteger(L, sockfd);
    lua_pushstring(L, addr);
    return 2;
}

int setnonblocking(int fd) {
    int flag = fcntl(fd, F_GETFL);
    if (fcntl(fd, F_SETFL, flag | O_NONBLOCK) < 0) {
        return -1;
    }
    return 0;
}

/*
    usage:  setnonblocking(fd)
    paras:  fd
    return: fd
*/
static int l_setnonblocking(lua_State *L) {
    int fd = luaL_checkinteger(L, 1);
    setnonblocking(fd);
    lua_pushinteger(L, fd);
    return 1;
}

int lua_lib_socket(lua_State *L) {
    static const struct luaL_Reg l[] = {
        // client
        { "connect", l_socket_connect },

        // server
        { "listen", l_socket_listen },
        { "accept", l_socket_accept },
        { "setnonblocking", l_setnonblocking },

        { "close", l_socket_close },
        { "send", l_socket_send },
        { "recv", l_socket_recv }, 
        { NULL, NULL }
    };
    luaL_newlib(L, l);
    return 1;
}